#! /usr/bin/env python
# -*- coding: utf-8 -*
import rospy
import numpy as np
from std_msgs.msg import String
from sensor_msgs.msg import Image, CameraInfo
import message_filters
import cv2 as cv
from cv_bridge import CvBridge, CvBridgeError
from geometry_msgs.msg import Twist
 

detect_sw = True
track_sw = False  
p1 = []
p = []
dst = []
num = 0


def robot_move(target,depth_image):
    z = 0
    x = 0
    target_x,target_y = target
    real_z = depth_image[int(target_y),int(target_x)] * 0.001 
    if target_x > 0 and target_y > 0 and target_x < 400 and target_y <640 and real_z > 0:
        
        msg = Twist()
        if real_z > 0:
            real_z = real_z-2.0
            if abs(real_z) > 0.6:
                if real_z>0:
                    z = 0.6
                else:
                    z = -0.6
            else:
                z = real_z
        else:
            z = 0
    
        x = (target[0]-320)/320*1.5
        if real_z > 0.5:
            real_z = 0.5
        print(z)
        msg.linear.x = z
        msg.angular.z = -x
        print("------------move----------------")
        pub_move.publish(msg)

def cascade_detection(cascade_file, image):
    cascade = cv.CascadeClassifier(cascade_file)
    gray = cv.cvtColor(image, cv.COLOR_BGR2GRAY)
    
    objs = cascade.detectMultiScale(gray,scaleFactor=1.1,minNeighbors=2,minSize=(50,50))
    result = False
    dst = []
    if len(objs) == 1:
        for obj in objs:    
            cv.rectangle(image,(obj[0],obj[1]),(obj[0]+obj[2],obj[1]+obj[3]),(255,0,0),2)
            cv.circle(image,(int(obj[0]+obj[2]/2),int(obj[1]+obj[3]*1/4)),10,(255,255,255),-1)
            dst =[int(obj[0]+obj[2]/2),int(obj[1]+obj[3]*1/4)]
            result = True
    return result,image,dst



def callback(data1,data2):
    global detect_sw,track_sw,p,dst,old_gray,num
    lk_params = dict(winSize = (30,30),maxLevel = 3, criteria = (cv.TERM_CRITERIA_EPS|cv.TERM_CRITERIA_COUNT,10,0.03))
    bridge = CvBridge()
    color_image = bridge.imgmsg_to_cv2(data1, 'bgr8')
    depth_image = bridge.imgmsg_to_cv2(data2, '16UC1')
    gray = cv.cvtColor(color_image,cv.COLOR_BGR2GRAY)
    cascade_file = rospy.get_param('~cascade_file','')
    # print(color_image.shape)
    
    if track_sw is True:
        print("目标跟随")
        detect_sw = False
        
        p1,st,err = cv.calcOpticalFlowPyrLK(old_gray,gray,p,None,**lk_params)
        num += 1
        p = p1

        cv.circle(color_image,(tuple(p.ravel())),5,(255,0,255),-1)
        print(p.ravel())
        robot_move(p.ravel(),depth_image)
        print(st)
        print(num)
        if st != 1:
            print("----------end-----------------")
            track_sw = False
            detect_sw = True
        if num > 50:
            num = 0
            track_sw = False
            detect_sw = True 
    

    if detect_sw is True:
        # print("检测中")
        pose = []
        detect_succeed,image,dst = cascade_detection(cascade_file,color_image)
        print(dst)
        if detect_succeed is True:
            # print("开始跟随")
            track_sw = True
            p = np.array([dst]).astype('float32').reshape(-1,1,2)
    old_gray = np.copy(gray)
    cv.imshow('color_image',color_image)
    cv.waitKey(10)
    


 
if __name__ == '__main__':
    global fx, fy, ppx, ppy
    fx = 414.8895263671875
    fy = 414.8895263671875
    ppx = 386.04364013671875
    ppy = 386.04364013671875

    rospy.init_node('get_image', anonymous=True)
 
    color = message_filters.Subscriber("/berxel_camera/rgb/rgb_raw", Image)
    depth = message_filters.Subscriber("/berxel_camera/depth/depth_raw", Image)
    pub_move = rospy.Publisher("/cmd_vel",Twist,queue_size=10)
    # color_depth = message_filters.ApproximateTimeSynchronizer([color, depth], 10, 1, allow_headerless=True)
    color_depth = message_filters.TimeSynchronizer([color, depth], 1)
    color_depth.registerCallback(callback)  
    
    rospy.spin()
    cv.destroyAllWindows()